package com.reebake.ideal.message.sse;

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class SseEmitterServer {

    private static final Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    /**
     * 创建连接
     */
    public SseEmitter connect(String userId) {
        if (sseEmitterMap.containsKey(userId)) {
            SseEmitter sseEmitter =sseEmitterMap.get(userId);
            sseEmitterMap.remove(userId);
            sseEmitter.complete();
        }
        try {

            // 设置超时时间，0表示不过期。默认30秒
            SseEmitter sseEmitter = new SseEmitter(0L);

            //给sseEmitter对象注册两个回调函数
            //当SseEmitter完成时执行的回调。这通常用于清理工作或记录日志。
            sseEmitter.onCompletion(completionCallBack(userId));
            //当SseEmitter因超时而关闭时执行的回调。这允许在连接超时时执行一些额外的操作
            sseEmitter.onTimeout(timeoutCallBack(userId));

            sseEmitterMap.put(userId, sseEmitter);
            log.debug("创建sse连接完成，当前用户：{}", userId);
            return sseEmitter;
        } catch (Exception e) {
            log.debug("创建sse连接异常，异常原因：{}", e.getMessage());
        }
        return null;
    }

    /**
     * 给指定用户发送消息
     * @param userId userId
     * @param messageId 消息标识
     * @param message 消息体，JSON格式
     * @return 布尔值
     */
    public boolean sendMessage(String userId,String messageId, String message) {
        if (sseEmitterMap.containsKey(userId)) {
            SseEmitter sseEmitter = sseEmitterMap.get(userId);
            try {
                sseEmitter.send(SseEmitter.event().id(messageId).data(message));
                log.debug("用户{},消息id:{},推送成功:{}", userId,messageId, message);
                return true;
            }catch (Exception e) {
                sseEmitterMap.remove(userId);
                log.debug("用户{},消息id:{},推送异常:{}", userId,messageId, e.getMessage());
                sseEmitter.complete();
                return false;
            }
        }else {
            log.debug("用户{}未上线", userId);
        }
        return false;
    }

    /**
     * 删除连接
     * @param userId userId
     */
    public void deleteUser(String userId){
        removeUser(userId);
    }

    private static Runnable completionCallBack(String userId) {
        return () -> {
            log.debug("结束sse用户连接：{}", userId);
            removeUser(userId);
        };
    }

    private static Throwable errorCallBack(String userId) {
        log.debug("sse用户连接异常：{}", userId);
        removeUser(userId);
        return new Throwable();
    }

    /**
     * 连接超时回调
     */
    private static Runnable timeoutCallBack(String userId) {
        return () -> {
            log.debug("连接sse用户超时：{}", userId);
            removeUser(userId);
        };
    }

    /**
     * 断开连接
     * @param userId userID
     */
    public static void removeUser(String userId){
        if (sseEmitterMap.containsKey(userId)) {
            SseEmitter sseEmitter = sseEmitterMap.get(userId);
            sseEmitterMap.remove(userId);
            //关闭与客户端的连接
            sseEmitter.complete();
        }else {
            log.debug("用户{} 连接已关闭，当前连接总数量：{}",userId,sseEmitterMap.size());
        }
    }

    public Map<String, SseEmitter> listSseConnect(){
        return sseEmitterMap;
    }
}

